use std::mem;
use windows_sys::Win32::System::{Diagnostics::Debug::{
    IMAGE_NT_HEADERS32, IMAGE_SECTION_HEADER,
}, SystemServices::IMAGE_DOS_HEADER};

#[allow(non_camel_case_types)]
#[allow(non_snake_case)]
#[repr(C)]
#[derive(Default)]
pub struct ANSI_STRING {
    Length: u16,
    MaximumLength: u16,
    Buffer: *mut i8,
}


fn image_first_section(nt_headers: *const IMAGE_NT_HEADERS32) -> *const IMAGE_SECTION_HEADER {
    unsafe {
        let ptr = nt_headers as *const u8;
        let offset = mem::offset_of!(IMAGE_NT_HEADERS32, OptionalHeader)
            + (*nt_headers).FileHeader.SizeOfOptionalHeader as usize;
        ptr.add(offset) as *const IMAGE_SECTION_HEADER
    }
}

fn get_section_from_rva(
    nt_headers: *const IMAGE_NT_HEADERS32,
    rva: u32,
) -> *const IMAGE_SECTION_HEADER {
    let mut section = image_first_section(nt_headers);
    unsafe {
        for _ in 0..(*nt_headers).FileHeader.NumberOfSections {
            if rva >= (*section).VirtualAddress
                && rva < (*section).VirtualAddress + (*section).SizeOfRawData
            {
                return section;
            }
            section = section.offset(1);
        }
    }
    std::ptr::null()
}

fn is_read_only_pointer(module_base: *const u8, pointer: *const u8) -> bool {
    if pointer.is_null() {
        return false;
    }

    let dos_header = module_base as *const IMAGE_DOS_HEADER;
    let nt_header = unsafe { module_base.add((*dos_header).e_lfanew as usize) as *const IMAGE_NT_HEADERS32 };

    let image_start = module_base as usize;
    let image_end = image_start + unsafe { (*nt_header).OptionalHeader.SizeOfImage as usize };
    let image_pointer = pointer as usize;

    if !(image_pointer > image_start && image_pointer < image_end) {
        return false;
    }

    let rva = (image_pointer - image_start) as u32;
    let section = get_section_from_rva(nt_header, rva);
    if section.is_null() {
        return false;
    }

    unsafe {
        let name_dword = *( &(*section).Name as *const [u8; 8] as *const u32 );
        let rdata_dword = u32::from_le_bytes([b'.', b'r', b'd', b'a']);
        name_dword == rdata_dword
    }
}

pub fn return_wow64_function_pointer(module_base: *const u8) -> *const u8 {

    let callback_function = ANSI_STRING {
        Length: 24,
        MaximumLength: 25,
        Buffer: b"Wow64PrepareForException\0".as_ptr() as *mut i8,
    };

    let dos_header = module_base as *const IMAGE_DOS_HEADER;
    let nt_header = unsafe { module_base.add((*dos_header).e_lfanew as usize) as *const IMAGE_NT_HEADERS32 };

    let mut section = image_first_section(nt_header);
    unsafe {
        for _ in 0..(*nt_header).FileHeader.NumberOfSections {
            let name_dword = *( &(*section).Name as *const [u8; 8] as *const u32 );
            let rdata_dword = u32::from_le_bytes([b'.', b'r', b'd', b'a']);
            if name_dword == rdata_dword {
                let rva = (*section).VirtualAddress;
                let entry_count = ((*section).Misc.VirtualSize - mem::size_of::<ANSI_STRING>() as u32)
                    / mem::size_of::<u64>() as u32;

                let pointer = (module_base.add(rva as usize)) as *const u64;

                for j in 0..entry_count {
                    let ptr_j = *pointer.add(j as usize);
                    if !is_read_only_pointer(module_base, ptr_j as *const u8) {
                        continue;
                    }

                    let api = ptr_j as *const ANSI_STRING;
                    if (*api).Length == callback_function.Length
                        && (*api).MaximumLength == callback_function.MaximumLength
                    {
                        let api_buf = std::slice::from_raw_parts(
                            (*api).Buffer as *const u8,
                            (*api).Length as usize,
                        );
                        let cb_buf = std::slice::from_raw_parts(
                            callback_function.Buffer as *const u8,
                            callback_function.Length as usize,
                        );
                        if api_buf == cb_buf {
                            return *pointer.add((j + 1) as usize) as *const u8;
                        }
                    }
                }
                break;
            }
            section = section.offset(1);
        }
    }

    std::ptr::null()
}

#[repr(C)]
pub struct FunctionAddressTable {
    pub nt_allocate_virtual_memory_address: u64,
    pub nt_protect_virtual_memory_address: u64,
    pub nt_write_virtual_memory: u64
}